package com.pai4j.zwiki.service;

import com.pai4j.zwiki.domain.entity.CodeStructureEntity;
import com.pai4j.zwiki.repository.CodeStructureRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import net.sourceforge.plantuml.FileFormat;
import net.sourceforge.plantuml.FileFormatOption;
import net.sourceforge.plantuml.SourceStringReader;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;

/**
 * PlantUML服务
 * 生成项目的UML架构图
 *
 * @author pai4j
 * @since 2025-10-12
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class PlantUmlService {

    private final CodeStructureRepository codeStructureRepository;

    @Value("${zwiki.workspace.base-path:/tmp/zwiki-workspace}")
    private String workspaceBasePath;

    /**
     * 生成类图
     *
     * @param taskId 任务ID
     * @return 类图文件路径
     */
    public String generateClassDiagram(String taskId) {
        log.info("开始生成类图, 任务ID: {}", taskId);

        // 获取代码结构
        List<CodeStructureEntity> structures = codeStructureRepository.findByTaskId(taskId);
        
        // 按包分组
        Map<String, List<CodeStructureEntity>> packageGroups = structures.stream()
                .filter(s -> "class".equals(s.getStructureType()) || "interface".equals(s.getStructureType()))
                .collect(Collectors.groupingBy(
                        s -> s.getPackageName() != null ? s.getPackageName() : "default"
                ));

        // 构建PlantUML类图代码
        StringBuilder uml = new StringBuilder();
        uml.append("@startuml\n");
        uml.append("' 项目类图\n");
        uml.append("skinparam classAttributeIconSize 0\n");
        uml.append("skinparam packageStyle rectangle\n\n");

        // 为每个包生成类
        for (Map.Entry<String, List<CodeStructureEntity>> entry : packageGroups.entrySet()) {
            String packageName = entry.getKey();
            List<CodeStructureEntity> classes = entry.getValue();

            if (!"default".equals(packageName)) {
                uml.append("package ").append(packageName).append(" {\n");
            }

            for (CodeStructureEntity clazz : classes) {
                String modifier = getUmlModifier(clazz.getAccessModifier());
                String type = "interface".equals(clazz.getStructureType()) ? "interface" : "class";
                
                uml.append("  ").append(modifier).append(" ").append(type)
                   .append(" ").append(clazz.getName()).append(" {\n");
                uml.append("  }\n");
            }

            if (!"default".equals(packageName)) {
                uml.append("}\n");
            }
            uml.append("\n");
        }

        uml.append("@enduml");

        // 生成图片
        return generateDiagram(taskId, "class-diagram", uml.toString());
    }

    /**
     * 生成架构图
     *
     * @param taskId 任务ID
     * @param architectureSummary 架构总结
     * @return 架构图文件路径
     */
    public String generateArchitectureDiagram(String taskId, String architectureSummary) {
        log.info("开始生成架构图, 任务ID: {}", taskId);

        // 获取包结构
        List<CodeStructureEntity> structures = codeStructureRepository.findByTaskId(taskId);
        Set<String> packages = structures.stream()
                .map(CodeStructureEntity::getPackageName)
                .filter(Objects::nonNull)
                .collect(Collectors.toSet());

        // 构建简化的架构图
        StringBuilder uml = new StringBuilder();
        uml.append("@startuml\n");
        uml.append("' 项目架构图\n");
        uml.append("!define RECTANGLE class\n\n");

        // 按包层级组织
        Map<String, List<String>> hierarchy = buildPackageHierarchy(packages);

        for (Map.Entry<String, List<String>> entry : hierarchy.entrySet()) {
            String parent = entry.getKey();
            List<String> children = entry.getValue();

            if (children.isEmpty()) {
                uml.append("component [").append(parent).append("]\n");
            } else {
                uml.append("package \"").append(parent).append("\" {\n");
                for (String child : children) {
                    uml.append("  [").append(child).append("]\n");
                }
                uml.append("}\n");
            }
        }

        uml.append("\n@enduml");

        return generateDiagram(taskId, "architecture-diagram", uml.toString());
    }

    /**
     * 生成组件图
     *
     * @param taskId 任务ID
     * @return 组件图文件路径
     */
    public String generateComponentDiagram(String taskId) {
        log.info("开始生成组件图, 任务ID: {}", taskId);

        List<CodeStructureEntity> structures = codeStructureRepository.findByTaskId(taskId);
        
        // 提取顶层包作为组件
        Set<String> topLevelPackages = structures.stream()
                .map(CodeStructureEntity::getPackageName)
                .filter(Objects::nonNull)
                .map(pkg -> {
                    String[] parts = pkg.split("\\.");
                    return parts.length >= 3 ? String.join(".", parts[0], parts[1], parts[2]) : pkg;
                })
                .collect(Collectors.toSet());

        StringBuilder uml = new StringBuilder();
        uml.append("@startuml\n");
        uml.append("' 项目组件图\n\n");

        for (String pkg : topLevelPackages) {
            String componentName = pkg.substring(pkg.lastIndexOf('.') + 1);
            uml.append("component [").append(componentName).append("] as ").append(componentName).append("\n");
        }

        uml.append("\n@enduml");

        return generateDiagram(taskId, "component-diagram", uml.toString());
    }

    /**
     * 生成PlantUML图表并保存为PNG
     */
    private String generateDiagram(String taskId, String diagramName, String plantUmlCode) {
        try {
            // 创建输出目录
            Path diagramDir = Paths.get(workspaceBasePath, taskId, "diagrams");
            Files.createDirectories(diagramDir);

            // 生成PNG文件
            String pngFilename = diagramName + ".png";
            Path pngFilePath = diagramDir.resolve(pngFilename);

            SourceStringReader reader = new SourceStringReader(plantUmlCode);
            try (FileOutputStream fos = new FileOutputStream(pngFilePath.toFile())) {
                reader.outputImage(fos, new FileFormatOption(FileFormat.PNG));
            }

            log.info("图表生成成功: {}", pngFilePath);
            return pngFilePath.toString();

        } catch (IOException e) {
            log.error("生成PlantUML图表失败", e);
            throw new RuntimeException("生成图表失败: " + e.getMessage(), e);
        }
    }

    /**
     * 获取UML访问修饰符
     */
    private String getUmlModifier(String accessModifier) {
        return switch (accessModifier) {
            case "public" -> "+";
            case "private" -> "-";
            case "protected" -> "#";
            default -> "~";
        };
    }

    /**
     * 构建包层级结构
     */
    private Map<String, List<String>> buildPackageHierarchy(Set<String> packages) {
        Map<String, List<String>> hierarchy = new HashMap<>();
        
        for (String pkg : packages) {
            String[] parts = pkg.split("\\.");
            if (parts.length > 3) {
                String parent = String.join(".", parts[0], parts[1], parts[2]);
                String child = parts[parts.length - 1];
                hierarchy.computeIfAbsent(parent, k -> new ArrayList<>()).add(child);
            } else {
                hierarchy.putIfAbsent(pkg, new ArrayList<>());
            }
        }
        
        return hierarchy;
    }
}

